import sys
import time
import numpy as np
import logging
from logging.handlers import RotatingFileHandler
import os
import torch
import torch.nn as nn
import ptan
import gymnasium as gym
from collections import deque, namedtuple
from torch.distributions import Categorical


HID_SIZE = 128


Experience = namedtuple('Experience', ['state', 'action', 'reward', 'done', 'next_state',  'ext_value', 'int_value', 'logprob', 'curiosity_reward'])


class RewardForwardFilter:
    def __init__(self, gamma):
        self.rewems = None
        self.gamma = gamma

    def update(self, rews):
        if self.rewems is None:
            self.rewems = rews
        else:
            self.rewems = self.rewems * self.gamma + rews
        return self.rewems


def _group_list(items, lens):
    """
    Unflat the list of items by lens
    反平铺队列，也就是将原先list 一维的数据，跟进lens进行分割，实现二维的队列
    [...] => [[.], [[.], [.]], .]
    :param items: list of items
    :param lens: list of integers
    :return: list of list of items grouped by lengths
    """
    res = []
    cur_ofs = 0
    for g_len in lens:
        res.append(items[cur_ofs:cur_ofs+g_len])
        cur_ofs += g_len
    return res


class RNDExperienceSource:

    def __init__(self, env, rnd_actor, rnd_model, steps_count=2, steps_delta=1, vectorized=False, device="cpu"):
        # 判断经验传入的参数类型是否正确
        # 并存储到成员变量
        assert isinstance(env, (gym.Env, list, tuple))
        assert isinstance(steps_count, int)
        assert steps_count >= 1
        assert isinstance(vectorized, bool)
        # self.pool: 存储游戏环境
        if isinstance(env, (list, tuple)):
            self.pool = env
        else:
            self.pool = [env]
        self.rnd_actor = rnd_actor
        self.rnd_model = rnd_model
        self.steps_count = steps_count
        self.steps_delta = steps_delta
        self.total_rewards = []
        self.total_steps = []
        self.vectorized = vectorized
        self.device = device

    def __iter__(self):
        states, agent_states, histories, cur_rewards, cur_steps = [], [], [], [], []
        # 存储每次环境观测值的长度（因为矢量的环境其返回的结果值是一个不定长度的列表）
        # 每个索引对应着states中对应索引的长度
        env_lens = []
        # 遍历每一个游戏环境
        # 这一大段的循环作用是初始化环境，得到初始化结果
        for env in self.pool:
            # 每一次遍历时，重置游戏环境
            obs, obs_info = env.reset()
            # if the environment is vectorized, all it's output is lists of results.
            # 如果支持矢量计算，那么可以将多个环境的输出结果，拼接到矩阵向量中，直接进行计算，效率比单个计算高
            if self.vectorized:
                # 矢量环境
                # 获取单词观察结果的向量长度
                obs_len = len(obs)
                # 将当前状态结果列表（应该是包含了环境状态，激励，动作，是否结束等信息）存储在states
                states.extend(obs)
            else:
                # 非矢量环境下，其观测的结果就是一个标量，简单说就是一个动作值，所以
                # 长度是1
                obs_len = 1
                # 将结果存储在status中
                states.append(obs)
            env_lens.append(obs_len)

            # 遍历本次环境观测的结果
            for _ in range(obs_len):
                histories.append(deque(maxlen=self.steps_count))  # 创建对应环境观测结果历史缓存队列
                cur_rewards.append(0.0)  # 存储最初是状态下的激励，为0
                cur_steps.append(0)  # 存储当前行走的部署，为0

        # 遍历索引
        # 从里这开始，应该就是尝试运行游戏了
        iter_idx = 0
        while True:
            actions = [None] * len(states)
            ext_values, int_values = [None] * len(states), [None] * len(states)
            logprobs = [None] * len(states)
            states_input = []
            states_indices = []
            # 遍历每一次的存储的观测状态
            # 对于非矢量环境来说，idx仅仅对应一个当前状态，但是对于矢量环境来说，idx对应当前获取的每个一观测值的索引
            # 这一个大循环的作用应该是，根据环境执行得到执行的动作结果
            # todo 有一个问题，矢量环境获取的多个观测结果这里怎么分辨
            for idx, state in enumerate(states):
                if state is None:
                    # 如果状态是空的，则使用环境进行随机选择一个动作执行， 另外这里假设的所有的环境都有相同的动作空间，
                    # 所以这里仅使用0索引环境进行随机动作采样
                    actions[idx] = self.pool[0].action_space.sample()  # assume that all envs are from the same family
                else:
                    # 如果状态非空，则将当前的存储在states环境观测值存储在states_input中
                    # 并存储当前的索引
                    states_input.append(state)
                    states_indices.append(idx)
            if states_input:
                # 如果观测的状态列表非空，则将状态输入的神经网络环境代理中，获取将要执行的动作
                # 而agent_staes根据源码，发现并未做处理
                states_v = ptan.agent.default_states_preprocessor(states_input).to(device=self.device)
                state_logits, _, _ = self.rnd_actor(states_v)
                probs = Categorical(logits=state_logits)
                state_actions = probs.sample()
                state_logprobs = probs.log_prob(state_actions)
                pred_ext_values, pred_int_values = self.rnd_actor.get_values(states_v)
                # 遍历每一个状态所要执行的动作
                for idx, (action, logprob, pred_ext_value, pred_int_value) in enumerate(zip(state_actions, state_logprobs.detach(), pred_ext_values.detach(), pred_int_values.detach())):
                    # 获取当前动作对应的状态的索引位置，有上面106行的代码可知
                    g_idx = states_indices[idx]
                    # 将执行的动作存储在与状态相对应的索引上
                    actions[g_idx] = action.cpu().numpy()
                    ext_values[g_idx] = pred_ext_value.flatten().cpu().numpy()
                    int_values[g_idx] = pred_int_value.flatten().cpu().numpy()
                    logprobs[g_idx] = logprob.cpu().numpy()
            # 将动作按照原先存储的每个环境得到的观测结果长度，按照环境数组进行重新分割分组
            grouped_actions = _group_list(actions, env_lens)
            grouped_ext_values = _group_list(ext_values, env_lens)
            grouped_int_values = _group_list(int_values, env_lens)
            grouped_logprobs = _group_list(logprobs, env_lens)

            # 因为存在一个大循环，存储每个环境的起始索引位置
            global_ofs = 0
            # 遍历每一个环境
            # 这一个大循环是将上一个循环中得到的执行动作应用到实际的环境中
            for env_idx, (env, action_n, ext_value_n, int_value_n, logprob_n) in enumerate(zip(self.pool, grouped_actions, grouped_ext_values, grouped_int_values, grouped_logprobs)):
                if self.vectorized:
                    # 这里action_n是一个list，也就是说矢量环境的输入的一个多维
                    # 如果是矢量的环境，则直接执行动作获取下一个状态，激励，是否结束等观测值
                    next_state_n, r_n, is_done_n, truncated, _ = env.step(action_n)
                    is_done_n = is_done_n or truncated
                else:
                    # 如果不是矢量环境，则需要将动作的第一个动作发送到env中获取相应的观测值（这里之所以是[0]，因为为了和矢量环境统一，即时是一个动作也会以列表的方式存储）
                    next_state, r, is_done, truncated, _ = env.step(action_n[0])
                    is_done = is_done or truncated
                    # 这个操作是为了和矢量环境统一
                    next_state_n, r_n, is_done_n = [next_state], [r], [is_done]

                # 遍历每一次的动作所得到的下一个状态、激励、是否结束
                for ofs, (action, next_state, r, is_done, ext_value, int_value, logprob) in enumerate(zip(action_n, next_state_n, r_n, is_done_n, ext_value_n, int_value_n, logprob_n)):
                    # 获取当前缓存的索引位置
                    idx = global_ofs + ofs
                    # 获取初始环境的状态
                    # 因为action_n存储的就是每一个状态下所执行的动作，所以这里直接使用idx提取对应的状态
                    state = states[idx]
                    # 获取一个历史队列，此时队列为空
                    history = histories[idx]

                    # 这里利用的idx来区分每一个状态执行的动作所对应的激励值
                    # 将获取的激励值存储在缓存中
                    cur_rewards[idx] += r
                    # 将当前状态以及执行的动作，执行的步骤次数存在起来
                    cur_steps[idx] += 1
                    next_state_v = ptan.agent.default_states_preprocessor([next_state]).to(device=self.device)
                    predict_next_feature, target_next_feature = self.rnd_model(next_state_v)
                    curiosity_reward = ((target_next_feature - predict_next_feature).pow(2).sum(1) / 2).detach().cpu().numpy()
                    # 如果状态非空，则（当前状态，所执行的动作对应的历史缓存队列）将当前状态存储在history中
                    # 所以一个样本可能就是这样对应一个队列数据
                    if state is not None:
                        history.append(Experience(state=state, action=action, reward=r, done=is_done, ext_value=ext_value, int_value=int_value, logprob=logprob, curiosity_reward=curiosity_reward, next_state=next_state))
                    # 如果达到了采集的步数并且遍历索引达到了两个经验样本的指定差值，则将样本返回，待外界下一次继续获取时，从这里继续执行
                    if len(history) == self.steps_count and iter_idx % self.steps_delta == 0:
                        yield tuple(history)
                    # 更新states，表示当前动作执行后状态的改变
                    # 将动作设置为动作执行后的下一个状态，因为idx表示当前运行的环境状态的变更
                    states[idx] = next_state
                    if is_done:
                        # 如果游戏结束，如果存储的历史数据小于指定的长度，则直接返回
                        # in case of very short episode (shorter than our steps count), send gathered history
                        if 0 < len(history) < self.steps_count:
                            yield tuple(history)
                        # generate tail of history
                        # 弹出最左侧的历史数据，返回给外部获取数据
                        while len(history) > 1:
                            history.popleft()
                            yield tuple(history)
                        # 将当前状态+动作执行后得到的激励存储在total_rewards队列中
                        self.total_rewards.append(cur_rewards[idx])
                        # 这个当前状态+动作执行的次数也存储起来
                        self.total_steps.append(cur_steps[idx])
                        # 重置状态
                        cur_rewards[idx] = 0.0
                        cur_steps[idx] = 0
                        # vectorized envs are reset automatically
                        states[idx] = env.reset()[0] if not self.vectorized else None
                        history.clear()

                # 将起始索引设置为下一个环境的起始位置
                global_ofs += len(action_n)
            # 遍历索引+1
            iter_idx += 1

    def pop_total_rewards(self):
        """
        返回所有采集的样本，并清空缓存
        """
        r = self.total_rewards
        if r:
            self.total_rewards = []
            self.total_steps = []
        return r

    def pop_rewards_steps(self):
        res = list(zip(self.total_rewards, self.total_steps))
        if res:
            self.total_rewards, self.total_steps = [], []
        return res


def setup_logger(save_path):
    logger = logging.getLogger(__name__)
    logger.setLevel(logging.INFO)
    handler = RotatingFileHandler(os.path.join(save_path, 'train.log'), maxBytes=1024 * 1024, backupCount=2)
    formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
    handler.setFormatter(formatter)
    logger.addHandler(handler)

    return logger


def save_best_model(score, state, save_dir, save_name, keep_best = 5):
    os.makedirs(save_dir, exist_ok=True)

    save_path = os.path.join(save_dir, f'{save_name}_{score}.pth')
    torch.save(state, save_path)

    all_model = sorted(filter(lambda x: "best" in x and "_" in x, os.listdir(save_dir)), key=lambda x: int(x.split('_')[-1].split('.')[0]))
    if len(all_model) > keep_best:
        for old_model in all_model[:-keep_best]:
            os.remove(os.path.join(save_dir, old_model))


class BipedalWalkerModelA2C(nn.Module):
    def __init__(self, obs_size, act_size):
        '''
        obs_size: 观测的数据维度，不是图像数据，所以后续用的是全连接层
        act_size: 动作空间的维度，在这个游戏里面，指的是同时执行动作的数量
        '''

        super(BipedalWalkerModelA2C, self).__init__()

        # 只有这个是提取特征，其余的都是输出结果
        self.base = nn.Sequential(
            nn.Linear(obs_size[0], HID_SIZE),
            nn.ReLU(),
        )
        # 输出均值，表示最终要执行的动作内容
        # 这里由于预测的输出动作包含负值，所以采用tanh函数，将输出值限制在-1到1之间
        # 而不是用sigmoid
        self.mu = nn.Sequential(
            nn.Linear(HID_SIZE, act_size[0]),
            nn.Tanh(),
        )
        # https: // zhuanlan.zhihu.com / p / 461707201
        # var作用 方差平方，怀疑是用来稳定输出的概率范围大小
        # 用来促进网络进行探索以及指导网路朝哪个方向进行训练，使得整体趋近中值，但是在这里并没有直接使用方差，而是使用了信息熵的方式
        self.var = nn.Sequential(
            nn.Linear(HID_SIZE, act_size[0]),
            nn.Softplus(), # Relu的替代函数，用于解决梯度消失问题 具体使用场景查看笔记内容
        )
        # 状态值（Q值），用来评价当前Q值，来评估当前执行的动作是否有优势
        self.value = nn.Linear(HID_SIZE, 1)

    def forward(self, x):
        '''
        return 均值，方差平方，Q值
        '''
        base_out = self.base(x)
        return self.mu(base_out), self.var(base_out), self.value(base_out)


class RewardTracker:
    def __init__(self, writer, stop_reward):
        self.writer = writer
        self.stop_reward = stop_reward

    def __enter__(self):
        self.ts = time.time()
        self.ts_frame = 0
        self.total_rewards = []
        return self

    def __exit__(self, *args):
        self.writer.close()

    def reward(self, reward, frame, epsilon=None):
        self.total_rewards.append(reward)
        speed = (frame - self.ts_frame) / (time.time() - self.ts)
        self.ts_frame = frame
        self.ts = time.time()
        mean_reward = np.mean(self.total_rewards[-100:])
        epsilon_str = "" if epsilon is None else ", eps %.2f" % epsilon
        print("%d: done %d games, mean reward %.3f, speed %.2f f/s%s" % (
            frame, len(self.total_rewards), mean_reward, speed, epsilon_str
        ))
        sys.stdout.flush()
        if epsilon is not None:
            self.writer.add_scalar("epsilon", epsilon, frame)
        self.writer.add_scalar("speed", speed, frame)
        self.writer.add_scalar("reward_100", mean_reward, frame)
        self.writer.add_scalar("reward", reward, frame)
        if mean_reward > self.stop_reward:
            print("Solved in %d frames!" % frame)
            return True
        return False


def unpack_batch_a2c(batch, net, last_val_gamma, device="cpu"):
    """
    Convert batch into training tensors
    :param batch: 收集的游戏数据
    :param net:
    :return: states variable, actions tensor, reference values variable（游戏环境状态、执行的动作、评价的Q值）
    """
    states = [] # 每一步的游戏状态
    actions = [] # 每一步执行的动作
    rewards = [] # 每一步执行动作后获取的奖励
    not_done_idx = [] # 执行动作后没有结束游戏的索引
    last_states = [] # 执行动作后的下一个状态
    for idx, exp in enumerate(batch):
        states.append(exp.state)
        actions.append(exp.action)
        rewards.append(exp.reward)
        if exp.last_state is not None:
            not_done_idx.append(idx)
            last_states.append(exp.last_state)
    states_v = ptan.agent.float32_preprocessor(states).to(device)
    actions_v = torch.FloatTensor(np.array(actions)).to(device)

    # handle rewards
    rewards_np = np.array(rewards, dtype=np.float32)
    if not_done_idx:
        # 如果存在下一个状态的游戏数据，那么计算Q值
        # 对于已经结束的游戏动作，他们的Q值
        last_states_v = ptan.agent.float32_preprocessor(last_states).to(device)
        last_vals_v = net(last_states_v)[2]
        # todo 为什么这里计算的Q值没有选择最大值
        last_vals_np = last_vals_v.data.cpu().numpy()[:, 0]
        rewards_np[not_done_idx] += last_val_gamma * last_vals_np

    ref_vals_v = torch.FloatTensor(rewards_np).to(device)
    return states_v, actions_v, ref_vals_v

def unpack_batch_ddqn(batch, device="cpu"):
    '''
    解压深度确定性策略梯度网络的数据
    '''
    # states: 环境状态
    # actions: 执行的动作
    # rewards： 执行动作后获取的奖励
    # dones: 执行动作后游戏是否结束
    # last_states: 未结束的游戏，执行动作后的达到的状态（针对多步展开，则是展开的最后一个动作）；如果是游戏已经结束的状态，则保存的还是和states中一样的状态，如果不是游戏结束的状态，则保存执行动作后的下一个状态
    states, actions, rewards, dones, last_states = [], [], [], [], []
    for exp in batch:
        states.append(exp.state)
        actions.append(exp.action)
        rewards.append(exp.reward)
        dones.append(exp.last_state is None)
        if exp.last_state is None:
            last_states.append(exp.state)
        else:
            last_states.append(exp.last_state)
    states_v = ptan.agent.float32_preprocessor(states).to(device)
    actions_v = ptan.agent.float32_preprocessor(actions).to(device)
    rewards_v = ptan.agent.float32_preprocessor(rewards).to(device)
    last_states_v = ptan.agent.float32_preprocessor(last_states).to(device)
    dones_t = torch.ByteTensor(dones).to(device)
    return states_v, actions_v, rewards_v, dones_t, last_states_v

def unpack_batch_acktr(batch, net, last_val_gamma, device="cpu"):
    """
    Convert batch into training tensors
    :param batch: 收集的游戏数据
    :param net:
    :return: states variable, actions tensor, reference values variable（游戏环境状态、执行的动作、评价的Q值）
    """
    states = [] # 每一步的游戏状态
    actions = [] # 每一步执行的动作
    rewards = [] # 每一步执行动作后获取的奖励
    not_done_idx = [] # 执行动作后没有结束游戏的索引
    last_states = [] # 执行动作后的下一个状态
    for idx, exp in enumerate(batch):
        states.append(exp.state)
        actions.append(exp.action)
        rewards.append(exp.reward)
        if exp.last_state is not None:
            not_done_idx.append(idx)
            last_states.append(exp.last_state)
    states_v = ptan.agent.float32_preprocessor(states).to(device)
    actions_v = torch.FloatTensor(np.array(actions)).to(device)

    # handle rewards
    rewards_np = np.array(rewards, dtype=np.float32)
    if not_done_idx:
        # 如果存在下一个状态的游戏数据，那么计算Q值
        # 对于已经结束的游戏动作，他们的Q值
        last_states_v = ptan.agent.float32_preprocessor(last_states).to(device)
        last_vals_v = net(last_states_v)
        # 为什么这里计算的Q值没有选择最大值
        # 因为评级网络是直接输出最大Q值，所以这里直接取值就可以了
        last_vals_np = last_vals_v.data.cpu().numpy()[:, 0]
        # 这里依旧是计算Q值
        rewards_np[not_done_idx] += last_val_gamma * last_vals_np

    ref_vals_v = torch.FloatTensor(rewards_np).to(device)
    return states_v, actions_v, ref_vals_v

def save_checkpoints(iter, state, checkpoint_dir, save_name, keep_last=5):
    if not os.path.exists(checkpoint_dir):
        os.makedirs(checkpoint_dir)

    checkpoint_path = os.path.join(checkpoint_dir, f'{save_name}_epoch_{iter}.pth')
    torch.save(state, checkpoint_path)

    all_checkpoints = sorted(filter(lambda x: "epoch" in x, os.listdir(checkpoint_dir)), key=lambda x: int(x.split('_')[2].split('.')[0]))
    if len(all_checkpoints) > keep_last:
        for old_checkpoint in all_checkpoints[:-keep_last]:
            os.remove(os.path.join(checkpoint_dir, old_checkpoint))